import numpy as np
import SimpleITK as sitk


def compute_dice_coefficient(mask_gt, mask_pred):
    """
    计算两个二值掩码之间的Dice系数

    参数:
    mask_gt (np.ndarray): 金标准分割掩码
    mask_pred (np.ndarray): 预测分割掩码

    返回:
    float: Dice系数值
    """
    # 计算交集体积和总体积
    volume_intersect = (mask_gt & mask_pred).sum()
    volume_sum = mask_gt.sum() + mask_pred.sum()
    # 计算并返回Dice系数
    return 2 * volume_intersect / volume_sum


# 设置金标准和预测结果的文件路径
gt_path = r"D:\zhuomian\BTCV\labelsTr\btca_0002.nii.gz"
pred_path = r"D:\zhuomian\BTCV\btca_0002_pred.nii.gz"

# 读取医学图像数据
gt_sitk = sitk.ReadImage(gt_path)
gt_data = sitk.GetArrayFromImage(gt_sitk)  # 转换为numpy数组
pred_sitk = sitk.ReadImage(pred_path)
pred_data = sitk.GetArrayFromImage(pred_sitk)  # 转换为numpy数组

# 遍历13个器官标签(1-13)，计算每个器官的Dice系数
for i in np.arange(1, 14):
    # 提取当前器官的二值掩码
    gt_i = gt_data == i
    pred_i = pred_data == i

    # 处理特殊情况：金标准和预测结果中都不存在该器官
    if np.sum(gt_i) == 0 and np.sum(pred_i) == 0:
        dsc = 1  # 完全匹配
    # 处理假阳性情况：金标准中不存在但预测存在
    elif np.sum(gt_i) == 0 and np.sum(pred_i) > 0:
        dsc = 0  # 完全不匹配
    # 处理假阴性情况：金标准中存在但预测不存在
    elif np.sum(gt_i) > 0 and np.sum(pred_i) == 0:
        dsc = 0  # 完全不匹配
    # 正常情况：计算实际Dice系数
    else:
        dsc = compute_dice_coefficient(gt_i, pred_i)

    # 输出每个器官的Dice系数
    print("标签 %d 的 Dice 系数为 %f" % (i, dsc))